from abc import ABCMeta, abstractmethod
import cv2
import torch
import numpy as np
from torch.utils.data import Dataset
from ..transform import Pipeline
import os


class BaseDataset(Dataset, metaclass=ABCMeta):
    """
    BaseDataset lika keras data  categorize with folders
    |-train
      |- cat
      |- car
      |- flower

    :param img_path: image data folder
    :param mode: train or val or test
    :param mode: 
    """

    def __init__(self,
                 img_path,
                 input_size,
                 pipeline,
                 keep_ratio=True,
                 class_name=None,
                 #  use_instance_mask=False,
                 #  use_seg_mask=False,
                 #  use_keypoint=False,
                 #  load_mosaic=False,
                 mode='train'
                 ):
        self.img_path = img_path
        self.input_size = input_size
        self.pipeline = Pipeline(pipeline, keep_ratio)
        self.keep_ratio = keep_ratio
        self.class_name = class_name
        # self.use_instance_mask = use_instance_mask
        # self.use_seg_mask = use_seg_mask
        # self.use_keypoint = use_keypoint
        # self.load_mosaic = load_mosaic
        self.mode = mode

        self.data_info = self.get_data_info(img_path)

    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, idx):
        if self.mode == 'val' or self.mode == 'test':
            return self.get_val_data(idx)
        else:
            while True:
                data = self.get_train_data(idx)
                if data is None:
                    idx = self.get_another_id()
                    continue
                return data

    def get_data_info(self, image_path):
        data_list = []
        for class_name in os.listdir(image_path):
            index = self.class_name.index(class_name)
            for filename in os.listdir(os.path.join(image_path,class_name)):
                file_path = os.path.join(image_path, class_name,filename)
                data_list.append([file_path, index])
        return data_list

    def get_train_data(self, idx):
        file_path, label = self.data_info[idx]
        # default use rgb image
        image = cv2.imdecode(np.fromfile(
            file_path, dtype=np.uint8), cv2.IMREAD_COLOR)[..., ::-1]
        width, height, channel = image.shape
        meta = dict(img=image,
                    img_info={'file_name': file_path,
                              'height': height,
                              'width': width,
                              'id': id}
                    )

        meta = self.pipeline(meta, self.input_size)
        meta['img'] = torch.tensor(meta['img'].transpose(2, 0, 1))
        # label = np.eye(len(self.class_name))[label]
        meta['label'] = torch.tensor(label)
        return meta

    def get_val_data(self, idx):
        return self.get_train_data(idx)

    def get_another_id(self):
        return np.random.random_integers(0, len(self.data_info)-1)
